import os
import re
import json
import requests
from datetime import datetime
from openai import OpenAI
import whisper
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_audio


AUDIO_DIR = 'audio'
JSON_DIR = 'json'
SUBTITLES_DIR = 'subtitles'
SUMMARY_DIR = 'summary'
TOOL_NAME = 'AI字幕工具 v0.0.2'
FORMAT_LIST = ['mp4', 'mpeg', 'mpga', 'm4a', 'wav', 'webm']
gpt_client = OpenAI(
    base_url = "https://oneapi.xty.app/v1",
    api_key = "sk-2izMV1pR78mIvUqg49Ed1aAbBa3746558292E2D5440cB6F7"
)


def get_duration(index):
    start_time = index * 5
    end_time = (index + 1) * 5
    return start_time, end_time


def time_to_seconds(time, time_format='%H:%M:%S,%f'):
    ref_dt = datetime(1900, 1, 1, 0, 0, 0)
    time = datetime.strptime(str(time), time_format) - ref_dt
    return time.seconds + time.microseconds / 1000000


def seconds_to_time(seconds, time_format='{:02d}:{:06.3f}'):
    return time_format.format(int(seconds // 60), seconds % 60)


def clear_folder(folder):
    if os.listdir(folder):
        if os.name == 'nt':
            os.system(f'cd /d "{folder}" & for /d %d in (*) do @rd /s /q "%d" & del /q /f "%d\*" 2>nul')
        else:
            os.system(f'rm -rf {folder}/*')


def replace_ext(filename, new_extension):
    root, _ = os.path.splitext(filename)
    new_filename = f'{root}.{new_extension}'
    return new_filename


def list_dir(folder, extensions=[]):

    def is_file_endswith(file, extensions):
        for extension in extensions:
            if file.endswith(extension):
                return True
        return False

    file_list = []
    for root, dirs, files in os.walk(folder):
        for file in files:
            if is_file_endswith(file, extensions):
                file_path = os.path.join(root, file)
                file_list.append(file_path)
    return file_list


def generate_srt(srt_path, subtitles):
    with open(srt_path, 'w') as fp:
        for i, subtitle in enumerate(subtitles, start=1):
            start_time = subtitle['start_time']
            end_time = subtitle['end_time']
            text = subtitle['text']
            srt = f"{i}\n{start_time} --> {end_time}\n{text}\n\n"
            fp.write(srt)


def download_bilibili(url):
    url = url.split('?')[0]
    print(f'Start downloading audio from {url}')
    session = requests.session()
    headers = {'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/106.0.0.0 Safari/537.36 Edg/106.0.1370.37',
            "Referer": "https://www.bilibili.com",
            }
    resp = session.get(url,headers=headers)
    title = re.findall(r'<title data-vue-meta="true">(.*?)_哔哩哔哩_bilibili',resp.text)[0]   
    play_info = re.findall(r'<script>window.__playinfo__=(.*?)</script>',resp.text)[0]

    json_data = json.loads(play_info) 
    audio_url = json_data['data']['dash']['audio'][0]['backupUrl'][0]  #音频地址  [0]清晰度最高
    video_url = json_data['data']['dash']['video'][0]['backupUrl'][0]   #视频地址

    audio_content = session.get(audio_url,headers=headers).content  #音频二进制内容
    video_content = session.get(video_url,headers=headers).content  #视频二进制内容

    print('Download complete.')
    return title, video_content, audio_content


def clean(save_dir=''):
    def delete_files_except_gitignore(folder_path):
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                os.remove(os.path.join(root, file))

    for folder in [AUDIO_DIR, JSON_DIR, SUBTITLES_DIR, SUMMARY_DIR]:
        print(f'cleaning up folder {folder} ...')
        folder_path = os.path.join(save_dir, folder)
        if os.path.exists(folder_path):
            delete_files_except_gitignore(folder_path)
            print('clean up done.')
        else:
            print(f"Folder '{folder}' does not exist.")

    print("Cleanup complete.")


def video2audio(video_list=[], save_dir=''):
    audio_dir = os.path.join(save_dir, AUDIO_DIR)
    if not os.path.exists(audio_dir):
        os.makedirs(audio_dir)
    for video_path in video_list:
        video_name = os.path.basename(video_path)
        print(f'converting {video_name} to audio ...')
        output_audio = os.path.join(audio_dir, replace_ext(video_name, 'mp3'))
        if not os.path.exists(output_audio):
            ffmpeg_extract_audio(video_path, output_audio)
        print(f'converting done.')
    print('converting to all videos to audios done.\n')


def audio2json(save_dir=''):
    # refer to: https://github.com/openai/whisper
    model = whisper.load_model('base')
    json_dir = os.path.join(save_dir, JSON_DIR)
    if not os.path.exists(json_dir):
        os.makedirs(json_dir)
    for audio_path in list_dir(os.path.join(save_dir, AUDIO_DIR), ['.mp3']):
        audio_name = os.path.basename(audio_path)
        print(f'extracting subtitles from {audio_name} ...')
        json_path = os.path.join(json_dir, replace_ext(audio_name, 'json'))
        json_result = model.transcribe(audio=audio_path, verbose=True, word_timestamps=True)
        with open(json_path, 'w') as fp:
            json.dump(json_result, fp)
        print('extracting done.')
    print('extracting all subtitles from audios done.\n')


def json2srt(save_dir=''):
    srt_dir = os.path.join(save_dir, SUBTITLES_DIR)
    if not os.path.exists(srt_dir):
        os.makedirs(srt_dir)
    for json_path in list_dir(os.path.join(save_dir, JSON_DIR), ['.json']):
        json_name = os.path.basename(json_path)
        print(f'generating srt file from {json_name} ...')
        srt_path = os.path.join(srt_dir, replace_ext(json_name, 'srt'))
        subtitles = []
        with open(json_path, 'r') as fp:
            data = json.load(fp)
            for segment in data['segments']:
                start_time = seconds_to_time(segment['start'])
                end_time = seconds_to_time(segment['end'])
                text = segment['text']
                subtitles.append({'start_time': start_time, 'end_time': end_time, 'text': text})
            generate_srt(srt_path, subtitles)
        print('generating done.')
    print('generating all srt files done.\n')


def get_completion(prompt, model="gpt-3.5-turbo", temperature=0):
    messages = [{"role": "user", "content": prompt}]
    try:
        response = gpt_client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
        )
        return response.choices[0].message.content

    except Exception as e:
        if hasattr(e, 'response') and e.response.status_code == 400:
            error_message = json.loads(e.response.text)['error']['message']
            print(error_message)


def summarize(save_dir=''):
    summary_dir = os.path.join(save_dir, SUMMARY_DIR)
    if not os.path.exists(summary_dir):
        os.makedirs(summary_dir)
    for srt_path in list_dir(os.path.join(save_dir, SUBTITLES_DIR), ['.srt']):
        srt_name = os.path.basename(srt_path)
        print(f'summarizing srt file {srt_name} ...')
        summary_path = os.path.join(summary_dir, replace_ext(srt_name, 'txt'))
        with open(summary_path, 'w') as fw:
            with open(srt_path, 'r') as fr:
                srt_text = fr.read()
                prompt = f'''
用简体中文尽可能详细地总结srt字幕文件内容，并提炼关键点，标出每个关键点的起始时间
输出格式如下：
```
总结性文字

[00:00.000] 关键点1
- [00:52.020] 相对于关键点1的子关键点1
- [08:26.460] 相对于关键点1的子关键点1

[08:77.460] 关键点2
- [09:52.020] 相对于关键点2的子关键点1
- [18:26.460] 相对于关键点2的子关键点2

......
```
需要总结的srt字幕文件内容如下：
```
{srt_text}
```
'''
                result = get_completion(prompt)
            if result:
                fw.write(result)
            else:
                print('summarize failed.')
                return
        print('summarizing done.')
    print('summarizing all srt files done.\n')


summarize(save_dir='/Users/mark/Downloads/xxx')